// Author: Ivan Chau
package uk.ac.cam.kcic2.photonMap;

import java.util.PriorityQueue;

import uk.ac.cam.kcic2.math.Params;
import uk.ac.cam.kcic2.math.Vector3D;
import uk.ac.cam.kcic2.scene.ColourCoeffs;
import uk.ac.cam.kcic2.scene.IntersectionInfo;

public class PhotonMap {

	private Photon[] photonMap;
	private int storedPhotons;

	public PhotonMap(int maxNumOfPhotons){
		photonMap = new Photon[maxNumOfPhotons];
		storedPhotons = 0;
	}

	// insert a photon into photon map
	public void insert(Photon p){
		photonMap[storedPhotons] = p;
		storedPhotons += 1;
	}

	public Photon get(int index){return photonMap[index];}

	public int size(){return storedPhotons;}

	public double[] getRadianceEstimates(IntersectionInfo intersection,
			boolean caustics, boolean mapView, double radius){
		if (storedPhotons == 0){
			return new double[]{0,0,0};
		}
		// locate n nearest photons
		Vector3D point = intersection.position;
		Vector3D normal = intersection.normal;
		ColourCoeffs coeffs = intersection.object.colourCoeffs;
		int n = Integer.MAX_VALUE;
		double k = 1.1;
		double r = radius;
		if (mapView){
			n = 1;
			r = 0.01;
		}
		Photon[] nearestPhotons = this.getKNN(n, point, r);
		// nearestPhotons are in the order of descending distance
		double[] sumOfFlux = new double[]{0,0,0};

		if (mapView){
			if (nearestPhotons.length > 0){
				assert nearestPhotons.length == 1;
				sumOfFlux[Params.X] = nearestPhotons[0].power.x;
				sumOfFlux[Params.Y] = nearestPhotons[0].power.y;
				sumOfFlux[Params.Z] = nearestPhotons[0].power.z;
			}
			return sumOfFlux;
		}

		for (Photon p : nearestPhotons){
			double cosTheta = p.direction.dot(normal);
			double d = Math.sqrt(p.position.sub(point).norm());
			if (cosTheta > 0){
				if (caustics){
					sumOfFlux[Params.X] += p.power.x*cosTheta*
							coeffs.diffuseCoeffs.x*
							Math.max(0, (1-d/(k*r)));
					sumOfFlux[Params.Y] += p.power.y*cosTheta
							*coeffs.diffuseCoeffs.y*
							Math.max(0, (1-d/(k*r)));
					sumOfFlux[Params.Z] += p.power.z*cosTheta*
							coeffs.diffuseCoeffs.z*
							Math.max(0, (1-d/(k*r)));
				}
				else{
					sumOfFlux[Params.X] += p.power.x*cosTheta*
							coeffs.diffuseCoeffs.x;
					sumOfFlux[Params.Y] += p.power.y*cosTheta*
							coeffs.diffuseCoeffs.y;
					sumOfFlux[Params.Z] += p.power.z*cosTheta*
							coeffs.diffuseCoeffs.z;
				}
			}
		}

		Vector3D radiance;

		if (caustics){
			radiance = new Vector3D(sumOfFlux).div((1-2/(3*k))*
					Math.PI*r*r);
		}
		else{
			radiance = new Vector3D(sumOfFlux).div(Math.PI*r*r);
		}
		return new double[]{radiance.x, radiance.y, radiance.z};
	}

	// turn the flat array into a heap like left-balanced kd-tree
	public void balanceTree(){
		Photon[] balancedPhotons = new Photon[storedPhotons];
		balance(balancedPhotons, photonMap, 0, 0, storedPhotons-1);
		photonMap = balancedPhotons;
	}

	// return the k nearest neighbours
	public Photon[] getKNN(int k, Vector3D queryPoint, double radius){		
		PhotonComparator pComp = new PhotonComparator(queryPoint);
		PriorityQueue<Photon> candidates = 
				new PriorityQueue<Photon>(1, pComp);
		findKNN(photonMap, 0, candidates, pComp, k, radius);
		Photon[] nearest = new Photon[candidates.size()];

		for (int i = 0; i < nearest.length; i++){
			nearest[i] = candidates.remove();
		}

		assert candidates.isEmpty();

		return nearest;
	}

	// insertion-sort the array in place based on the given axis
	public void insertionSort(Photon[] photons, int start, int end, 
			short plane){
		for (;start < end; start++){
			// find the min between start and end inclusively
			int min = start;
			for (int i = start+1; i <= end; i++){
				if (photons[min].getPosition((int)plane) > 
				photons[i].getPosition((int)plane)){
					min = i;
				}
			}
			// swap the elements
			swap(photons, start, min);
		}
		return;
	}

	// swap photons[a] and photons[b]
	private void swap(Photon[] photons, int a, int b){
		Photon temp = photons[a];
		photons[a] = photons[b];
		photons[b] = temp;
		return;
	}

	// find the median of medians of the array
	private int findMedianOfMedians(Photon[] photons, int start, 
			int end, short plane){
		int numOfNodes = end-start+1;

		if (numOfNodes < 6){	// just insertion-sort and return
			insertionSort(photons, start, end, plane);
			int median = start + ((numOfNodes+1)>>1)-1;
			return median;
		}
		else{
			int numOfGroups = Math.round(numOfNodes/5);
			int remainder = numOfNodes%5;

			// divide into groups of 5 and find the medians
			if (remainder == 0){
				for (int group = 0; group < numOfGroups; group ++){
					insertionSort(photons, start+group*5,
							start+group*5+4, plane);
				}

				// swap all the medians to the front
				for (int group = 0; group < numOfGroups; group ++){
					swap(photons, start+group, start+group*5+2);
				}				
			}
			else{
				for (int group = 0; group < numOfGroups-1; group ++){
					insertionSort(photons, start+group*5, 
							start+group*5+4, plane);
				}
				// for the last group
				insertionSort(photons, start+(numOfGroups-1)*5, end, 
						plane);

				// swap all the medians to the front
				for (int group = 0; group < numOfGroups-1; group ++){
					swap(photons, start+group, start+group*5+2);
				}
				// for the last group
				swap(photons, start+numOfGroups-1, 
						start+(numOfGroups-1)*5+((remainder+1)>>1)-1);
			}

			// find the median of the medians recursively
			return findMedianOfMedians(photons, start, 
					start+numOfGroups-1, plane);
		}
	}

	// partition the array around the pivot and return the index of 
	// the pivot
	private int partition(Photon[] photons, int start, int end, 
			short plane, int pivot){
		// put the pivot at the end
		swap(photons, pivot, end);

		int endOfLeftPartition = start-1;
		boolean leftFlag = true;
		for (int index = start; index < end; index++){
			if (photons[index].getPosition(plane) < 
					photons[end].getPosition(plane)){
				// put the photon into the left partition
				endOfLeftPartition += 1;
				swap(photons, endOfLeftPartition, index);
			}
			else if (photons[index].getPosition(plane) == 
					photons[end].getPosition(plane)){
				// put the photon into each partition alternatively
				if (leftFlag){
					// put the photon into the left partition
					endOfLeftPartition += 1;
					swap(photons, endOfLeftPartition, index);
					// flip the flag
					leftFlag = false;
				}
				else{
					// flip the flag
					leftFlag = true;
				}
			}
		}
		// put the pivot in place
		swap(photons, endOfLeftPartition+1, end);

		return endOfLeftPartition+1;
	}

	// find the ith smallest element using the Median of Medians
	// Algorithm
	public Photon findIthSmallest(Photon[] photons, int start, int end,
			short plane, int i){
		int numOfNodes = end-start+1;

		if (numOfNodes < 6){	// just insertion-sort and return
			insertionSort(photons, start, end, plane);
			return photons[i];
		}

		// find the median of medians
		int medianOfMedians = findMedianOfMedians(photons, start, end,
				plane);
		// partition around the median of medians
		medianOfMedians = partition(photons, start, end, plane, 
				medianOfMedians);

		if (i == medianOfMedians){
			return photons[i];
		}
		else if (i < medianOfMedians){
			// look at the left partition
			return findIthSmallest(photons, start, medianOfMedians-1,
					plane, i);
		}
		else{
			// look at the right partition
			return findIthSmallest(photons, medianOfMedians+1, end,
					plane, i);
		}
	}


	private void balance(Photon[] balancedPhotons, Photon[] photons,
			int index, int start, int end){		
		int numOfNodes = end-start+1;
		// if the partition contains none or one element
		if (numOfNodes <= 1){
			if (numOfNodes == 1){
				balancedPhotons[index] = photons[start];
			}
			return;
		}

		// compute the bounding box of the partition
		double[] bBoxMax = new double[]{Double.NEGATIVE_INFINITY, 
				Double.NEGATIVE_INFINITY, Double.NEGATIVE_INFINITY};
		double[] bBoxMin = new double[]{Double.POSITIVE_INFINITY, 
				Double.POSITIVE_INFINITY, Double.POSITIVE_INFINITY};
		for (int photonIndex = start; photonIndex <= end; 
				photonIndex++){
			for (int spaceIndex = Params.X; 
					spaceIndex < Params.VECTOR3D_DIM; spaceIndex++){
				bBoxMax[spaceIndex] = Math.max(bBoxMax[spaceIndex], 
						photons[photonIndex].getPosition(spaceIndex));
				bBoxMin[spaceIndex] = Math.min(bBoxMin[spaceIndex], 
						photons[photonIndex].getPosition(spaceIndex));
			}
		}

		// choose the axis as the comparator where the photons differ
		// the most
		short plane = Params.Z;
		if (bBoxMax[Params.X] - bBoxMin[Params.X] > 
		bBoxMax[Params.Y] - bBoxMin[Params.Y] && 
		bBoxMax[Params.X] - bBoxMin[Params.X] > 
		bBoxMax[Params.Z] - bBoxMin[Params.Z]){
			plane = Params.X;
		}
		else if (bBoxMax[Params.Y] - bBoxMin[Params.Y] > 
		bBoxMax[Params.Z] - bBoxMin[Params.Z]){
			plane = Params.Y;
		}

		// compute the desired ith smallest element, such that 
		// choosing this as the root will give a left balanced tree

		int depth = Integer.SIZE - Integer.numberOfLeadingZeros(
				numOfNodes) - 1;
		int maxNumOfNodes = (int)(Math.pow(2,(depth+1))) - 1;
		int maxNumOfLeaves = (int) Math.pow(2, depth);
		int numOfMissingLeaves = maxNumOfNodes - numOfNodes;
		int pivot;
		if (numOfMissingLeaves*2 > maxNumOfLeaves){	
			// a non-full left subtree
			// the right subtree is full with depth one smaller than
			// the left subtree
			pivot = end - ((int)(Math.pow(2, depth-1)) - 1);
		}
		else{	// a full left subtree
			pivot = start + ((int)(Math.pow(2, depth)) - 1);
		}

		// find the pivot element in linear time using Median of 
		// Medians Algorithm
		// the array is partitioned around the pivot as a side effect
		Photon pivotPhoton = findIthSmallest(photons, start, end, 
				plane, pivot);
		pivotPhoton.setPlane(plane);
		balancedPhotons[index] = pivotPhoton;
		// recursively find the pivot element in the partitions to put
		// the array into a heap
		if (2*index+1 < storedPhotons){	// if a left subtree exists
			balance(balancedPhotons, photons, 2*index+1, start, 
					pivot-1);
		}
		if (2*index+2 < storedPhotons){	// if a right subtree exists
			balance(balancedPhotons, photons, 2*index+2, pivot+1, end);
		}
		return;
	}

	// find the k nearest neighbours if any recursively
	private void findKNN(Photon[] photons, int root, 
			PriorityQueue<Photon> candidates, PhotonComparator pComp,
			int k, double radius){
		assert candidates.size() <= k;

		short plane = photons[root].getPlane();
		// if not leaf
		if (2*root+1 < storedPhotons){
			double distance = pComp.queryPoint.get(plane) - 
					photons[root].getPosition(plane);

			if (distance > 0){	// search the right plane first
				if (2*root+2 < storedPhotons){	
					// if a right child exists
					findKNN(photons, 2*root+2, candidates, pComp, k, 
							radius);
				}
				if (distance < radius){	
					// the left plane may contain candidates
					findKNN(photons, 2*root+1, candidates, pComp, k, 
							radius);
				}
			}
			else{	// search the left plane first
				findKNN(photons, 2*root+1, candidates, pComp, k, 
						radius);
				if (-distance < radius){ 
					// the right plane may contain candidates
					if (2*root+2 < storedPhotons){ 
						// if a right child exists
						findKNN(photons, 2*root+2, candidates, pComp,
								k, radius);
					}
				}
			}
		}

		double distSqrd = photons[root].position.sub(
				pComp.queryPoint).norm();

		if (distSqrd < radius*radius){
			// a photon is found
			if (candidates.size() < k){ // the queue is not full
				candidates.add(photons[root]);
			}
			else{ // check if this photon is a better candidate
				if (pComp.compare(
						photons[root], candidates.peek()) > 0){
					candidates.remove();
					candidates.add(photons[root]);
				}
			}
		}

		assert candidates.size() <= k;
		return;
	}

}
